use std::sync::Arc;

use futures::StreamExt;
use polars_core::frame::DataFrame;
use polars_error::PolarsResult;
use polars_plan::dsl::sink2::FileProviderArgs;
use polars_utils::IdxSize;

use crate::async_executor::{self, TaskPriority};
use crate::nodes::io_sinks2::components::error_capture::ErrorCapture;
use crate::nodes::io_sinks2::components::file_sink::FileSinkPermit;
use crate::nodes::io_sinks2::components::hstack_columns::HStackColumns;
use crate::nodes::io_sinks2::components::partition_sink_starter::PartitionSinkStarter;
use crate::nodes::io_sinks2::components::partition_state::PartitionState;
use crate::nodes::io_sinks2::components::sink_morsel::SinkMorsel;
use crate::nodes::io_sinks2::components::size::RowCountAndSize;

pub struct PartitionMorselSender {
    /// Note: Must be <= `file_size_limit` if there is one.
    pub ideal_morsel_size: RowCountAndSize,
    pub file_size_limit: Option<RowCountAndSize>,
    pub inflight_morsel_semaphore: Arc<tokio::sync::Semaphore>,
    pub open_sinks_semaphore: Arc<tokio::sync::Semaphore>,
    pub partition_sink_starter: PartitionSinkStarter,
    /// For include_key: true
    pub hstack_keys: Option<HStackColumns>,
    pub error_capture: ErrorCapture,
}

pub enum NextMorsel {
    TakeFromBuffered {
        num_rows: IdxSize,
    },
    /// Sorted finalize will receive morsels from a stream.
    Morsel(SinkMorsel),
}

impl PartitionMorselSender {
    /// # Panics
    /// Panics if `partition.file_sink_task_data` is `None`.
    pub async fn send_morsels(
        &self,
        partition: &mut PartitionState,
        flush: bool,
        // Sorted finalize
        morsel_stream: Option<(futures::stream::BoxStream<'static, (SinkMorsel, bool)>, u64)>,
    ) -> PolarsResult<()> {
        let row_byte_size: u64;

        let mut sorted_finalize_stream = if let Some((stream, row_byte_size_)) = morsel_stream {
            row_byte_size = row_byte_size_;

            assert_eq!(partition.sinked_size, RowCountAndSize::default());
            assert_eq!(partition.buffered_rows.height(), 0);

            Some(stream)
        } else {
            row_byte_size = partition.buffered_size().row_byte_size();
            None
        };

        loop {
            let file_sink_task_data = partition.file_sink_task_data.as_mut().unwrap();

            let mut used_row_capacity: RowCountAndSize;
            let mut available_row_capacity: RowCountAndSize;

            macro_rules! calc_used_and_available_capacities {
                ($file_sink_task_data:expr) => {{
                    used_row_capacity = partition
                        .sinked_size
                        .checked_sub($file_sink_task_data.start_position)
                        .unwrap();
                    available_row_capacity =
                        self.file_size_limit
                            .map_or(RowCountAndSize::MAX, |file_size_limit| RowCountAndSize {
                                num_rows: file_size_limit
                                    .num_rows
                                    .checked_sub(used_row_capacity.num_rows)
                                    .unwrap(),
                                num_bytes: file_size_limit
                                    .num_bytes
                                    .saturating_sub(used_row_capacity.num_bytes),
                            });
                }};
            }

            calc_used_and_available_capacities!(file_sink_task_data);

            if used_row_capacity.num_rows == 0 && available_row_capacity.num_rows == 0 {
                available_row_capacity = RowCountAndSize {
                    num_rows: 1,
                    num_bytes: u64::MAX,
                };
            }

            let (next_morsel, start_new_sink) = if let Some(sorted_finalize_stream) =
                sorted_finalize_stream.as_mut()
            {
                let Some((morsel, start_new_sink)) = sorted_finalize_stream.next().await else {
                    return Ok(());
                };

                (NextMorsel::Morsel(morsel), start_new_sink)
            } else {
                let buffered_size = partition.buffered_size();

                if buffered_size.num_rows == 0 {
                    return Ok(());
                }

                let num_rows_to_take = self.ideal_morsel_size.num_rows_takeable_from(buffered_size);

                let could_buffer_more = num_rows_to_take == buffered_size.num_rows
                    && num_rows_to_take < self.ideal_morsel_size.num_rows;

                if could_buffer_more && !flush {
                    return Ok(());
                }

                let max_takeable_rows: IdxSize =
                    available_row_capacity.num_rows_takeable_from(buffered_size);

                let start_new_sink = max_takeable_rows == 0;
                let num_rows_to_take = if start_new_sink {
                    num_rows_to_take
                } else {
                    num_rows_to_take.min(max_takeable_rows)
                };

                (
                    NextMorsel::TakeFromBuffered {
                        num_rows: num_rows_to_take,
                    },
                    start_new_sink,
                )
            };

            if start_new_sink {
                assert!(used_row_capacity.num_rows > 0);
                let handle = partition.file_sink_task_data.take().unwrap().close();

                let file_permit: FileSinkPermit =
                    if let Ok(permit) = self.open_sinks_semaphore.clone().try_acquire_owned() {
                        async_executor::spawn(
                            TaskPriority::Low,
                            self.error_capture.clone().wrap_future(handle),
                        );

                        permit
                    } else {
                        handle.await?
                    };

                partition.file_sink_task_data = Some(self.partition_sink_starter.start_sink(
                    FileProviderArgs {
                        index_in_partition: partition.num_sink_opens,
                        partition_keys: partition.keys_df.clone(),
                    },
                    partition.sinked_size,
                    file_permit,
                )?);
                partition.num_sink_opens += 1;
            }

            let file_sink_task_data = partition.file_sink_task_data.as_mut().unwrap();

            if start_new_sink {
                calc_used_and_available_capacities!(file_sink_task_data);
            }

            let mut morsel = match next_morsel {
                NextMorsel::TakeFromBuffered { num_rows } => {
                    let (df, remaining) = partition.buffered_rows.split_at(
                        #[allow(clippy::unnecessary_fallible_conversions)]
                        i64::try_from(num_rows).unwrap(),
                    );

                    partition.buffered_rows = remaining;

                    let morsel_permit = self
                        .inflight_morsel_semaphore
                        .clone()
                        .acquire_owned()
                        .await
                        .unwrap();

                    SinkMorsel::new(df, morsel_permit)
                },
                NextMorsel::Morsel(sink_morsel) => sink_morsel,
            };

            let morsel_height: IdxSize = IdxSize::try_from(morsel.df().height()).unwrap();

            debug_assert!(
                self.ideal_morsel_size.num_rows
                    <= self.file_size_limit.map_or(IdxSize::MAX, |x| x.num_rows)
            );
            debug_assert!(morsel_height <= self.ideal_morsel_size.num_rows);

            assert!((1..=available_row_capacity.num_rows).contains(&morsel_height));

            if let Some(hstack_keys) = self.hstack_keys.as_ref() {
                let columns = morsel.df().get_columns();
                let height = morsel.df().height();
                let new_columns = hstack_keys.hstack_columns_broadcast(
                    height,
                    columns,
                    partition.keys_df.get_columns(),
                );

                *morsel.df_mut() = unsafe { DataFrame::new_no_checks(height, new_columns) };
            };

            if file_sink_task_data.morsel_tx.send(morsel).await.is_err() {
                let handle = partition.file_sink_task_data.take().unwrap().close();
                return Err(handle.await.unwrap_err());
            }

            let sinked_size_delta = RowCountAndSize {
                num_rows: morsel_height,
                #[allow(clippy::useless_conversion)]
                num_bytes: u64::from(morsel_height)
                    .saturating_mul(row_byte_size)
                    .min(available_row_capacity.num_bytes),
            };

            partition.sinked_size = partition
                .sinked_size
                .checked_add(sinked_size_delta)
                .unwrap();
        }
    }
}
